from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, count

from SparkSessionBase import SparkSessionBase


class BusinessApi(SparkSessionBase):
    SPARK_URL = "yarn"
    SPARK_APP_NAME = 'business_name_select'
    ENABLE_HIVE_SUPPORT = True

    def __init__(self):
        self.spark = self._create_spark_session()
        self.spark.sparkContext.setLogLevel("ERROR")  # 设置日志级别为 ERROR，减少日志

    def get_business_stats(self, business_id):
        # 从 business 表中获取商户信息
        business_df = self.spark.table('business')
        business_info = business_df.filter(col("business_id") == business_id).first()

        if not business_info:
            return None  # 如果没有找到商户，返回 None

        # 获取 business 表中的所有字段
        business_data = business_info.asDict()

        # 从 review 表中获取评论统计数据
        review_df = self.spark.table('review')
        review_stats = review_df.filter(col("rev_business_id") == business_id).agg(
            sum("rev_stars").alias("total_rev_stars"),
            sum("rev_useful").alias("total_rev_useful"),
            sum("rev_funny").alias("total_rev_funny"),
            sum("rev_cool").alias("total_rev_cool")
        )

        # 从 clean_tips 表中获取 tips 统计数据
        tips_df = self.spark.table('clean_tip')
        tips_count = tips_df.filter(col("business_id") == business_id).count()

        # 将 business 数据与统计结果合并
        stats_data = {
            **business_data,
            "total_rev_stars": review_stats.first()["total_rev_stars"],
            "total_rev_useful": review_stats.first()["total_rev_useful"],
            "total_rev_funny": review_stats.first()["total_rev_funny"],
            "total_rev_cool": review_stats.first()["total_rev_cool"],
            "total_tips": tips_count
        }

        # 创建一个 DataFrame 来存储统计结果
        stats_df = self.spark.createDataFrame([stats_data])

        # 返回统计结果 DataFrame
        return stats_df

    def run(self, business_id):
        stats_df = self.get_business_stats(business_id)
        if stats_df:
            stats_df.show(truncate=False)  # 显示统计结果
            print("business success")
            return stats_df
        else:
            print(f"No business found with ID: {business_id}")
            return None


if __name__ == '__main__':
    # 示例调用
    job = BusinessApi()
    job.run(business_id="8rb-3VYXE37IZix4yOdskw")